Python Pytorch实现图像识别

发表于 | 次阅读
字数统计: 1,843 字 | 阅读时长 ≈ 9 min

环境准备

1.python 环境，建议python3.8

环境搭建

pip install torch torchvision torchaudio

数据集准备

新建 pytorch_test 目录,并在目录下 新建RGB 文件夹及子目录，如下

1
2
3
4
5
6
7
8
9
└─RGB
    ├─train
    │  ├─green
    │  ├─red
    │  └─yellow
    └─validation
        ├─green
        ├─red
        └─yellow

在green，red,yellow 文件中分别放好，相应底色的图片，每个文件夹，图片数量都需要大于1

数据集验证

查看数据集代码
data_read.py

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
import torch.utils.data
import numpy as np
import os, random, glob
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt


# 数据集读取
class MyDataSet(torch.utils.data.Dataset):
    def __init__(self, img_dir, transform=None):
        self.transform = transform 
        g_dir = os.path.join(img_dir, "green")  #####要训练的.bmp上一级的名字，这里设你要区分的类别为N，这里的N=3.
        r_dir = os.path.join(img_dir, "red")    #####
        y_dir = os.path.join(img_dir, "yellow")  ##### 

        imgsLib = [] 
        imgsLib.extend(glob.glob(os.path.join(g_dir, "*.bmp")))   #####将图片的地址添加到列表
        imgsLib.extend(glob.glob(os.path.join(r_dir, "*.bmp")))   #####
        imgsLib.extend(glob.glob(os.path.join(y_dir, "*.bmp")))   #####
        random.shuffle(imgsLib)  # 打乱数据集
        self.imgsLib = imgsLib

    # 作为迭代器必须要有的
    def __getitem__(self, index):
        img_path = self.imgsLib[index]

        label = 1 if 'green' in img_path.split('/')[-1] else 0                    #####为类别做标签，三类就012
        if label == 0:  label = 2 if 'yellow' in img_path.split('/')[-1] else 0   #####

        img = Image.open(img_path).convert("RGB")
        img = self.transform(img)
        return img, label

    def __len__(self):
        return len(self.imgsLib)

# 数据集读取
# class MyDataSet2(torch.utils.data.Dataset):
 
#  读取数据

if __name__ == "__main__":

    CLASSES = {0: "red", 1: "green", 2: "yellow"}   ###### 标号与类别名做好对应
    #img_dir = "C:/Users/ASUS/Desktop/project/RGB/train" #####地址为你自己train的位置，最好是绝对寻址
    img_dir = "./pytorch_test/RGB/train" #####地址为你自己train的位置，最好是绝对寻址
    data_transform = transforms.Compose([
        transforms.Resize(256),  # resize到256   #####规定尺寸，直到运行这部分代码可以完全看见图片
        transforms.CenterCrop(512),  # crop到224 #####
        transforms.ToTensor(),
    ])

    dataSet = MyDataSet(img_dir=img_dir, transform=data_transform)
    # dataLoader = torch.utils.data.DataLoader(dataSet, batch_size=8, shuffle=True, num_workers=4)
    # dataLoader = torch.utils.data.DataLoader(dataSet, batch_size=8, shuffle=False, num_workers=4)
    dataLoader = torch.utils.data.DataLoader(dataSet, batch_size=8, shuffle=False, num_workers=1)


    #image_batch, label_batch = iter(dataLoader).next()
    #image_batch, label_batch = next(dataLoader)
    #image_batch, label_batch = next(iter(dataLoader))

    image_batch, label_batch = next(iter(dataLoader))

    for i in range(image_batch.data.shape[0]):
        label = np.array(label_batch.data[i])          
        # print(label)
        img = np.array(image_batch.data[i]*255, np.int32)
        print(CLASSES[int(label)])
        plt.imshow(np.transpose(img, [1, 2, 0]))
        plt.show()

运行查看，是否正确读取数据集

报错处理

问题一

1
ValueError: num_samples should be a positive integer value, but got num_samples=0

代码调试时，显示ValueError: num_samples should be a positive integer value, but got num_samples=0
因为用的数据集是已经划分好的，所以不需要再shuffle。加载数据时将shuffle = False，错误即可消除。

1
# dataLoader = torch.utils.data.DataLoader(dataSet, batch_size=8, shuffle=True, num_workers=4)

修改为

1
2
3
dataLoader = torch.utils.data.DataLoader(dataSet, batch_size=8, shuffle=False, num_workers=4)
``` 
#### 问题二

AttributeError: ‘_MultiProcessingDataLoaderIter’ object has no attribute ‘next’

1
2
```
image_batch, label_batch = iter(dataLoader).next()

运行上述语句，报错：AttributeError: ‘_MultiProcessingDataLoaderIter’ object has no attribute ‘next’

先尝试将其改为单线程处理

1
2
3
4
5
6
7
8
9
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                        shuffle=True，num_workers=0)
```                         
                                        
发现问题并没有解决，又继续报错，AttributeError: '_SingleProcessingDataLoaderIter' object has no attribute 'next'

出现原因：pytorch版本关于next()的用法不一样

解决方法：

image_batch, label_batch = next(iter(dataLoader))

1
2
3
4
5
 即可成功运行
 

## 模型训练
训练代如下

from future import print_function, division
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data import Dataset
from torchvision import transforms, datasets, models

from data_read1 import MyDataSet #####第一个.py文件名

配置参数

random_state = 1
torch.manual_seed(random_state) # 设置随机数种子，确保结果可重复
torch.cuda.manual_seed(random_state)
torch.cuda.manual_seed_all(random_state)
np.random.seed(random_state)

random.seed(random_state)

epochs = 20 # 训练次数
batch_size = 16 # 批处理大小
num_workers = 0 # 多线程的数目
use_gpu = torch.cuda.is_available()

对加载的图像作归一化处理， 并裁剪为[224x224x3]大小的图像

data_transform = transforms.Compose([
transforms.Resize(256), #####这部分尺寸最好一致
transforms.CenterCrop(512), #####
transforms.ToTensor(),

])

train_dataset = MyDataSet(img_dir=”./pytorch_test/RGB/train”, transform=data_transform) #####训练数据集的地址

#train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=False, num_workers=0)

test_dataset = MyDataSet(img_dir=”./pytorch_test2/RGB/validation”, transform=data_transform)#####测试数据集的地址

#test_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
test_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=False, num_workers=0)

加载resnet18 模型，

net = models.resnet18(pretrained=False)
num_ftrs = net.fc.in_features
net.fc = nn.Linear(num_ftrs, 3) # 更新resnet18模型的fc模型，#####因为这里有三类，所以设置为3

if use_gpu:
net = net.cuda()
print(net)

‘’’
Net (
(conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1))
(maxpool): MaxPool2d (size=(2, 2), stride=(2, 2), dilation=(1, 1))
(conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
(fc1): Linear (44944 -> 2048)
(fc2): Linear (2048 -> 512)
(fc3): Linear (512 -> 2)
)
‘’’

定义loss和optimizer

cirterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.0001, momentum=0.9)

开始训练

net.train()
if name == ‘main‘:
for epoch in range(epochs):
running_loss = 0.0
train_correct = 0
train_total = 0
for i, data in enumerate(train_loader, 0):
inputs, train_labels = data
if use_gpu:
inputs, labels = Variable(inputs.cuda()), Variable(train_labels.cuda())
else:
inputs, labels = Variable(inputs), Variable(train_labels)
# inputs, labels = Variable(inputs), Variable(train_labels)
optimizer.zero_grad()
outputs = net(inputs)
_, train_predicted = torch.max(outputs.data, 1)
# import pdb
# pdb.set_trace()
train_correct += (train_predicted == labels.data).sum()
loss = cirterion(outputs, labels)
loss.backward()
optimizer.step()

        running_loss += loss.item()
        print("epoch: ",  epoch, " loss: ", loss.item())
        train_total += train_labels.size(0)

    print('train %d epoch loss: %.3f  acc: %.3f ' % (
    epoch + 1, running_loss / train_total * batch_size, 100 * train_correct / train_total))

    # 模型测试
    correct = 0
    test_loss = 0.0
    test_total = 0
    test_total = 0
    net.eval()
    for data in test_loader:
        images, labels = data
        if use_gpu:
            images, labels = Variable(images.cuda()), Variable(labels.cuda())
        else:
            images, labels = Variable(images), Variable(labels)
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        loss = cirterion(outputs, labels)
        test_loss += loss.item()
        test_total += labels.size(0)
        correct += (predicted == labels.data).sum()

    print('test  %d epoch loss: %.3f  acc: %.3f ' % (epoch + 1, test_loss / test_total, 100 * correct / test_total))

torch.save(net, “./pytorch_test/my_model1.pth”) #####保存为一个my_model1.pth文件

1
2
3
4
5

运行之后，可在文件夹下找到my_model1.th 模型文件

## 验证模型
验证代码

from PIL import Image
import torch
from torchvision import transforms

图片路径

#save_path = ‘./my_model3.pth’ #####上一次训练生成的文件
save_path = ‘./pytorch_test/my_model1.pth’ #####上一次训练生成的文件

———————— 加载数据 —————————

Data augmentation and normalization for training

Just normalization for validation

定义预训练变换

preprocess_transform = transforms.Compose([
transforms.Resize(256), #####尺寸保持一致
transforms.CenterCrop(512), #####
transforms.ToTensor(), #####
# transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

class_names = [‘red’, ‘green’, ‘yellow’] #####类别的名字，按顺序

device = torch.device(“cuda:0” if torch.cuda.is_available() else “cpu”)

———————— 载入模型并且测试 —————————

model = torch.load(save_path)
model.eval()

print(model)

#image_PIL = Image.open(‘./pytorch_test/RGB/validation/green/1.bmp’) #####放入你要测试的图片位置
image_PIL = Image.open(‘./pytorch_test/RGB/validation/red/2.bmp’) #####放入你要测试的图片位置
#image_PIL = Image.open(‘./pytorch_test/RGB/validation/yellow/3.bmp’) #####放入你要测试的图片位置

image_tensor = preprocess_transform(image_PIL)

以下语句等效于 image_tensor = torch.unsqueeze(image_tensor, 0)

image_tensor.unsqueeze_(0)

没有这句话会报错

image_tensor = image_tensor.to(device)

out = model(image_tensor)

得到预测结果，并且从大到小排序

_, indices = torch.sort(out, descending=True)

返回每个预测值的百分数

percentage = torch.nn.functional.softmax(out, dim=1)[0]

print(out)
print(“图片预测为：{}”.format(class_names[int(out.argmax(1))]))

print([(class_names[idx], percentage[idx].item()) for idx in indices[0][:5]])

1
2

输出预测结果

tensor([[ 0.0585, -0.0298, -0.0243]], grad_fn=)
图片预测为：red
[(‘red’, 0.35260626673698425), (‘yellow’, 0.32458174228668213), (‘green’, 0.322812020778656)]

```

一辈子很短，努力的做好两件事就好；
第一件事是热爱生活，好好的去爱身边的人；
第二件事是努力学习，在工作中取得不一样的成绩，实现自己的价值，而不是仅仅为了赚钱；

继开 wechat
欢迎加我的微信，共同交流技术